﻿// Copyright (c) Microsoft. All rights reserved.

using System.Collections.Generic;

#pragma warning disable IDE0130 // reduce number of "using" statements
// ReSharper disable once CheckNamespace - reduce number of "using" statements
namespace Microsoft.KernelMemory;

/// <summary>
/// Settings used by the default SearchClient
/// </summary>
public class SearchClientConfig
{
    /// <summary>
    /// Maximum number of tokens accepted by the LLM used to generate answers.
    /// The number includes the tokens used for the answer, e.g. when using
    /// GPT4-32k, set this number to 32768.
    /// If the value is not set or less than one, SearchClient will use the
    /// max amount of tokens supported by the model in use.
    /// </summary>
    public int MaxAskPromptSize { get; set; } = -1;

    /// <summary>
    /// Maximum number of relevant sources to consider when generating an answer.
    /// The value is also used as the max number of results returned by SearchAsync
    /// when passing a limit less or equal to zero.
    /// </summary>
    public int MaxMatchesCount { get; set; } = 100;

    /// <summary>
    /// How many tokens to reserve for the answer generated by the LLM.
    /// E.g. if the LLM supports max 4000 tokens, and AnswerTokens is 300, then
    /// the prompt sent to LLM will contain max 3700 tokens, composed by
    /// prompt + question + grounding information retrieved from memory.
    /// </summary>
    public int AnswerTokens { get; set; } = 300;

    /// <summary>
    /// Text to return when the LLM cannot produce an answer.
    /// </summary>
    public string EmptyAnswer { get; set; } = "INFO NOT FOUND";

    /// <summary>
    /// Template use to inject facts into the RAG prompt.
    /// Available placeholders:
    /// * {{$content}}  : text from memory, i.e. chunk of text extracted from the source
    /// * {{$source}}   : name of the source file, or URL of the web page, where the content originated.
    /// * {{$relevance}}: relevance score of the current chunk of text
    /// * {{$memoryId}} : ID of the memory record
    /// * {{$tags}}     : list of tags, excluding reserved/internal ones
    /// * {{$tag[X]}}   : tag X value(s), replaced with "-" if the value is empty
    /// * {{$meta[X]}}  : value of memory record payload X field (memory payload is also known as metadata), replaced with "-" if the value is empty
    /// </summary>
    public string FactTemplate { get; set; } = "==== [File:{{$source}};Relevance:{{$relevance}}]:\n{{$content}}";

    /// <summary>
    /// The memory DB might include duplicate chunks of text, e.g. when importing the same files
    /// with different document IDs or chat messages (high probability), or when the same text
    /// appears in different files (not very frequent, considering partitioning process).
    /// If two chunks are equal (not case-sensitive), regardless of tags and file names, it's usually
    /// better to skip the duplication, including the chunk only once in the RAG prompt, reducing the
    /// tokens used. The chunk will still be listed under sources.
    /// You might want to set this to True if your prompt includes other chunk details, such as tags
    /// and filenames, that could affect the LLM output.
    /// Note: when the storage contains duplicate records, other relevant records will be left out,
    /// possibly affecting RAG quality, because deduplication occurs after retrieving N records from storage,
    /// leaving RAG with [N - count(duplicates)] records to work with.
    /// </summary>
    public bool IncludeDuplicateFacts { get; set; } = false;

    /// <summary>
    /// Number between 0.0 and 2.0. It controls the randomness of the completion.
    /// The higher the temperature, the more random the completion.
    /// </summary>
    public double Temperature { get; set; } = 0;

    /// <summary>
    /// Number between 0.0 and 2.0. It controls the diversity of the completion.
    /// The higher the TopP, the more diverse the completion.
    /// </summary>
    public double TopP { get; set; } = 0;

    /// <summary>
    /// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether
    /// they appear in the text so far, increasing the model's likelihood to talk about
    /// new topics.
    /// </summary>
    public double PresencePenalty { get; set; } = 0;

    /// <summary>
    /// Number between -2.0 and 2.0. Positive values penalize new tokens based on their
    /// existing frequency in the text so far, decreasing the model's likelihood to repeat
    /// the same line verbatim.
    /// </summary>
    public double FrequencyPenalty { get; set; } = 0;

    /// <summary>
    /// Up to 4 sequences where the completion will stop generating further tokens.
    /// </summary>
    public IList<string> StopSequences { get; set; } = [];

    /// <summary>
    /// Modify the likelihood of specified tokens appearing in the completion.
    /// </summary>
    public Dictionary<int, float> TokenSelectionBiases { get; set; } = [];

    /// <summary>
    /// Whether to check is the generated answers are safe.
    /// A content moderation service must be present in the system.
    /// </summary>
    public bool UseContentModeration { get; set; } = true;

    /// <summary>
    /// Answer to return when AI generates content considered unsafe by the
    /// text moderation service.
    /// </summary>
    public string ModeratedAnswer { get; set; } = "Sorry, the generated content contains unsafe or inappropriate information.";

    /// <summary>
    /// Verify that the current state is valid.
    /// </summary>
    public void Validate()
    {
        if (this.MaxAskPromptSize is > 0 and < 1024)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.MaxAskPromptSize)} cannot be less than 1024");
        }

        if (this.MaxMatchesCount < 1)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.MaxMatchesCount)} cannot be less than 1");
        }

        if (this.AnswerTokens < 1)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.AnswerTokens)} cannot be less than 1");
        }

        if (this.EmptyAnswer.Length > 256)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.EmptyAnswer)} is too long, consider something shorter");
        }

        if (this.Temperature is < 0 or > 2)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.Temperature)} must be between 0 and 2");
        }

        if (this.TopP is < 0 or > 2)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.TopP)} must be between 0 and 2");
        }

        if (this.PresencePenalty is < -2 or > 2)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.PresencePenalty)} must be between -2 and 2");
        }

        if (this.FrequencyPenalty is < -2 or > 2)
        {
            throw new ConfigurationException($"SearchClient: {nameof(this.FrequencyPenalty)} must be between -2 and 2");
        }
    }
}
